import math
import numpy as np

import torch

from einops import rearrange, reduce, repeat


@torch.no_grad()
def association_knn(
    token_positions: torch.Tensor, 
    batch_token_ids, 
    super_token_positions: torch.Tensor,
    batch_super_token_ids,
    k: int,
):
    """
    associate super token and token by searching k nearest neighbor \
    super tokens for each token
    
    Args:
        token_positions (torch.Tensor): [T, D] in voxel / pixel coordinates
        batch_token_ids (list[torch.Tensor]): token indices for each batch 
        super_token_positions (torch.Tensor): [S, D]
        batch_super_token_ids (list[torch.Tensor]): super token indices for each batch 
        k: number of nearest neighbors selected for each token
        
    Returns:
        linked_pairs (torch.Tensor): [T * K, 2] token-supertoken index pairs
        num_linked_super_tokens (torch.Tensor): [T]
        num_linked_tokens (torch.Tensor): [S]
    """
    # constants
    T = token_positions.shape[0]
    S = super_token_positions.shape[0]
    k = min(k, S)
    device = token_positions.device
    # compute relations
    linked_pairs = []
    for t_inds, st_inds in zip(batch_token_ids, batch_super_token_ids):
        t_pos = token_positions[t_inds].float()
        st_pos = super_token_positions[st_inds].float()
        dist = torch.cdist(t_pos, st_pos, p=2)
        knn_st_inds = dist.topk(k, dim=-1, largest=False).indices
        knn_st_inds = st_inds[rearrange(knn_st_inds, "t k -> (t k)")]
        knn_t_inds = repeat(t_inds, "t -> (t k)", k=k)
        linked_pairs.append(torch.stack([knn_t_inds, knn_st_inds], dim=-1))
    linked_pairs = torch.cat(linked_pairs, dim=0)
    # compute num linked
    num_linked_super_tokens = torch.zeros(T, device=device, dtype=torch.long)
    num_linked_super_tokens.index_add_(0, linked_pairs[:, 0], torch.ones_like(linked_pairs[:, 0]))
    num_linked_tokens = torch.zeros(S, device=device, dtype=torch.long)
    num_linked_tokens.index_add_(0, linked_pairs[:, 1], torch.ones_like(linked_pairs[:, 1]))
    return linked_pairs, num_linked_super_tokens, num_linked_tokens


ASSOCIATION_FNS = {
    "knn": association_knn,
}